﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using MotionPlanning.MotionPlanningPrimitives;
using MotionPlanning.MotionPlanningComponents;
using MotionPlanning.MotionPlanningComponents.Visualizations;
using MotionPlanning.MotionPlanningComponents.MotionPlanningSystems;

namespace MotionPlanning.MotionPlanningAlgorithms.RRTas_Standard
{
    class RRTas : Planner_RRTBased
    {

        public RRTas(RoboticSystemComponent system, VisualizationComponent visualizer, string motionPlanningConfigFile)
            : base(system, visualizer, motionPlanningConfigFile)
        {
        }

        /// <summary>
        /// Execute RRT\ast-with-shooting algorithm
        /// </summary>
        /// <param name="n_iterations"></param>
        public override void Execute(int n_iterations)
        {
            if (x_root == null)
                throw new Exception("x_root uninitialized. Forget to call SetRoot()?");
            int iterations_thisRound = 0;
            while (iterations_thisRound < n_iterations)
            {
                ++iterations_total;
                ++iterations_thisRound;
                #region SharedVariables
                Node x_nearest = null;
                Node x_new = null;
                #endregion
                #region Extension Phase
                bool DoExtension = true;
                if (DoExtension)
                {
                    /*
                    ----------------- Sampling Phase -------------------
                    */
                    double[] q_sample = this.Sampler.NextQ(SamplingComponent.SamplingMode.SAMPLING, this.System);
                    /*
                     ---------- Nearest Neighbor Search Phase -----------
                     */
                    double dist_nearest = double.MaxValue;
                    Object obj_nearest = null;
                    this.Vertices.NearestNeighbor_Approximate(q_sample, ref dist_nearest, ref obj_nearest);
                    if (obj_nearest == null) continue;
                    x_nearest = (Node)obj_nearest;
                    /*
                     ----------------- Extension Phase ------------------
                     */
                    //propagation
                    double[] u_inc = this.Sampler.NextControl(SamplingComponent.SamplingMode.CONTROL, this.System);
                    double t_inc = u_inc[u_inc.Length - 1];
                    double[] q_new = new double[this.System.Get_DOF_state()];
                    double cost_inc = 0;
                    bool succ_propagation = this.System.Propagate(x_nearest.get_Value().get_q(), u_inc, t_inc, ref q_new, ref cost_inc);
                    if (succ_propagation == false) continue;
                    bool inRange = this.System.PointInRange(q_new);
                    if (inRange == false) continue;
                    //grow tree
                    x_new = this.SpawnNode_As(q_new);
                    this.Vertices.AddPoint(x_new.get_Value().get_q(), x_new);

                    Edge e_new = this.SpawnEdge(x_nearest, x_new);
                    e_new.get_Value().set_ctrl_inc(u_inc);
                    e_new.get_Value().set_utitlity_inc_c(cost_inc);

                    x_new.get_Value().set_utility_c(x_nearest.get_Value().get_utility_c() + cost_inc);

                    x_nearest.add_ChildEdge(e_new);
                    x_new.add_ParentEdge(e_new);
                }
                #endregion
                #region Improve Phase

                bool DoImprovement = true;
                if (DoImprovement)
                {
                    int n_shoot = 100;

                    for (int i = 0; i < n_shoot; ++i)
                    {
                        //shooting
                        double[] u_inc = this.Sampler.NextControl(SamplingComponent.SamplingMode.SHOOTING, this.System);
                            //modify u_inc[u_inc.Length - 1] if needed
                        double t_inc = u_inc[u_inc.Length - 1]; ///TODO change this hand coded value
                        double[] q_close = new double[this.System.Get_DOF_state()];
                        double cost_inc_shoot = 0;
                        bool succ_propagation = this.System.Propagate_HighSpeed(x_nearest.get_Value().get_q(), u_inc, t_inc, ref q_close, ref cost_inc_shoot);
                        if (succ_propagation == false) continue;
                        //matching
                        double delta_close = 0.01; ///TODO put this into a class member
                        double dist_nearest2 = double.MaxValue;
                        Object obj_nearest = null;
                        this.Vertices.NearestNeighbor_Approximate(q_close, ref dist_nearest2, ref obj_nearest);
                        if (obj_nearest == null) continue;
                        Node x_near = (Node)obj_nearest;
                        //skip if shooting deteriorates to find x_nearest
                        if (x_near == x_nearest)
                            continue;
                        //skip if gap is too large
                        if (dist_nearest2 > delta_close * delta_close)
                            continue;
                        //skip if shooting is not benefitial
                        double cost_at_x_close = x_nearest.get_Value().get_utility_c() + cost_inc_shoot;
                        if (cost_at_x_close >= x_near.get_Value().get_utility_c())
                            continue;
                        //Rewiring
                        Edge e_shoot = this.SpawnEdge(x_nearest, x_near);
                        e_shoot.get_Value().set_utitlity_inc_c(cost_inc_shoot);
                        e_shoot.get_Value().set_ctrl_inc(u_inc);
                        List<Edge> E_parent = x_near.get_Parents();
                        Edge e_parent = E_parent[0];
                        Node x_parent = e_parent.get_entry();
                        x_parent.get_Children().Remove(e_parent);
                        E_parent.Remove(e_parent);
                        e_parent.set_entry(null);
                        e_parent.set_exit(null);
                        e_parent.set_Value(null);
                        x_near.add_ParentEdge(e_shoot);
                        x_near.get_Value().set_utility_c(cost_at_x_close);
                        double[] q_near =  x_near.get_Value().get_q();
                        Vertices.DelPoint(q_near, x_near);
                        MotionPlanningUtils.vector_set(q_close, ref q_near);
                        Vertices.AddPoint(q_near, x_near);
                        x_nearest.add_ChildEdge(e_shoot);
                        //propagate the changes in the sub-branchs
                        branchSize = 0;
                        Traverse(x_near, 0);
                        if (branchSize > 100000) continue;
                        SwingBranch(x_near);
                    }
                }
                #endregion
            }
        }

        private int branchSize;
        private void Traverse(Node x_src, int level)
        {
            branchSize++;
            List<Edge> E_children = x_src.get_Children();
            if (E_children == null) return;
            for (int i = 0; i < E_children.Count; ++i)
            {
                Edge e_i = E_children[i];
                Node x_brother = e_i.get_exit();
                Traverse(x_brother, level + 1);
            }
        }

        private void SwingBranch(Node x)
        {
            List<Edge> E_children = x.get_Children();
            if (E_children == null) return;
            double cost_inc = 0;
            double[] q_tmp = new double[this.System.Get_DOF_state()];
            for (int i = 0; i < E_children.Count; ++i)
            {
                Edge e_i = E_children[i];
                double[] u_inc = e_i.get_Value().get_ctrl_inc();
                double t_inc = u_inc[u_inc.Length - 1];
                Node x_i = e_i.get_exit();
                bool succ_propagation = this.System.Propagate(x.get_Value().get_q(), u_inc, t_inc, ref q_tmp, ref cost_inc);
                if (succ_propagation == false) continue;
                double[] q_i = x_i.get_Value().get_q();
                this.Vertices.DelPoint(q_i, x_i);
                MotionPlanningUtils.vector_set(q_tmp, ref q_i);
                this.Vertices.AddPoint(q_i, x_i);
                e_i.get_Value().set_utitlity_inc_c(cost_inc);
                x_i.get_Value().set_utility_c(x.get_Value().get_utility_c() + cost_inc);
                
                SwingBranch(x_i);
            }
        }

        public override Node SpawnNode_Like(double[] q)
        {
            int DOF_state = this.System.Get_DOF_state();
            double[] q_new = new double[DOF_state];
            for (int i = 0; i < DOF_state; ++i)
                q_new[i] = q[i];
            Node x = Node.SpawnGraphicNode();
            x.set_Value(new NodeData_RRTas(q_new, 0));
            return x;
        }

        public override Node SpawnNode_As(double[] q)
        {
            if (q.Length != this.System.Get_DOF_state())
                throw new Exception("q.length != DOF_state. in RRTas.SpawnNode_As");
            Node x = Node.SpawnGraphicNode();
            x.set_Value(new NodeData_RRTas(q, 0));
            return x;
        }
        public override Edge SpawnEdge()
        {
            Edge e = Edge.SpawnGraphicEdge();
            e.set_Value(new EdgeData_RRTas());
            return e;
        }
        public override Edge SpawnEdge(Node entry, Node exit)
        {
            Edge e = Edge.SpawnGraphicEdge(entry, exit);
            e.set_Value(new EdgeData_RRTas());
            return e;
        }
        public override string ToString()
        {
            throw new NotImplementedException();
        }
        public override void Test()
        {
            throw new NotImplementedException();
        }

        /// <summary>
        ///  test if q_test is in the goal region
        /// </summary>
        /// <param name="q_test"></param>
        /// <returns></returns>
        public override bool GoalReached(double[] q_test)
        {
            if (r_goal <= 0)
                throw new NotImplementedException("Forget to call SetGoal and SetGoalTolerance first?");
            double distance2 = this.System.Distance2_Euclidean(q_goal, q_test);
            if (distance2 <= r_goal * r_goal)
                return true;
            else
                return false;
        }
        public double GetGoalTolerance()
        {
            return this.r_goal;
        }
        public void SetGoalTolerance(double radius)
        {
            this.r_goal = radius;
        }
        override public void SetGoal(double[] q_goal)
        {
            this.q_goal = q_goal;
        }
        public double[] GetGoal()
        {
            return q_goal;
        }

        override public void Plot()
        {
            List<object> graph = Vertices.GetAllDataObjects();
            Visualizer.PlotVerticesWithCost(graph);
        }
    }
}
